# ***************将处理好的.txt文件分为训练集以及验证集，并且放到固定路径中************
import os
import random
from PIL import Image
from tqdm import tqdm
import shutil

# random.seed(42)

train = "D://PyCharm/Code/yolov5-5.0/yolov5-5.0/data/JEPGImages/"
label = "D://PyCharm/Code/yolov5-5.0/yolov5-5.0/data/labels/"

save_images_train = "D://PyCharm/Code/yolov5-5.0/yolov5-5.0/VOCdevkit/images/train/"
save_images_val = "D://PyCharm/Code/yolov5-5.0/yolov5-5.0/VOCdevkit/images/val/"

save_labels_train = "D://PyCharm/Code/yolov5-5.0/yolov5-5.0/VOCdevkit/labels/train/"
save_labels_val = "D://PyCharm/Code/yolov5-5.0/yolov5-5.0/VOCdevkit/labels/val/"

num = 0
img_paths = []
xml_paths = []

images_train = []
labels_trian = []
images_val = []
labels_val = []

for fname in os.listdir(train):
    img_path = os.path.join(train, fname)
    xml_path = os.path.join(label, fname.split('.')[0] + '.txt')
    year = int(fname.split('_')[0])
    if year == 2021:
        img_paths.append(img_path)
        xml_paths.append(xml_path)
        # print("{}, {}".format(img_path, xml_path))
val = int(len(img_paths) * 0.1)  # 百分之10的训练集
train = int(len(img_paths) - val)

for i in range(val):
    flag = random.randint(0, len(img_paths)) - 1
    # label = random.randint(0, val)
    images_val.append(img_paths[flag])
    labels_val.append(xml_paths[flag])
    # print("{}, {}".format(img_paths[flag], xml_paths[flag]))

images_train = [x for x in img_paths if x not in images_val]
labels_trian = [y for y in xml_paths if y not in labels_val]


# print("{}, {}, {}, {}".format(len(images_val), len(labels_val), len(images_train), len(labels_trian)))

def save_images(lit, save_dir):
    inters = tqdm(range(len(lit)))
    for i in inters:
        img = Image.open(lit[i])
        img = img.convert('RGB')
        name = lit[i].split('/')
        save = save_dir + name[-1]
        img.save(save)
    inters.update(1)


def save_labels(lit, save_dir):
    inters = tqdm(range(len(lit)))
    # file = open(save_dir)
    for i in inters:
        xml_oldname = lit[i]
        # img = img.convert('RGB')
        # file = open(lit[i])
        name = lit[i].split('/')
        # print(name)
        xml_newname = save_dir + name[-1]
        shutil.copyfile(xml_oldname, xml_newname)
    inters.update(1)


if __name__ == '__main__':
    save_images(images_train, save_images_train)
    save_labels(labels_trian, save_labels_train)
    save_images(images_val, save_images_val)
    save_labels(labels_val, save_labels_val)
